import caffe
import caffe.proto.caffe_pb2
import numpy as np
from caffe.io import blobproto_to_array
import os

#load the model
net = caffe.Net('bvlc_alexnet/deploy.prototxt',
                'alexnet/caffe_alexnet_train_iter_23627.caffemodel',
                caffe.TEST)

blob = caffe.proto.caffe_pb2.BlobProto()
data = open( 'data_mean.binaryproto' , 'rb' ).read()
blob.ParseFromString(data)
arr = np.array( caffe.io.blobproto_to_array(blob) )


# load input and configure preprocessing
transformer = caffe.io.Transformer({'data': net.blobs['data'].data.shape})
transformer.set_mean('data', arr[0].mean(1).mean(1))
transformer.set_transpose('data', (2,0,1))
transformer.set_channel_swap('data', (2,1,0))
transformer.set_raw_scale('data', 255.0)

#note we can change the batch size on-the-fly
#since we classify only one image, we change batch size from 10 to 1
# net.blobs['data'].reshape(1,3,85,85)

#load the image in the data layer

filelist = open('data.test/filelist.txt')
lines = filelist.readlines()
for line in lines:
    line = line.split(' ')
    imgname = line[0]
    label = int(line[1])
    im = caffe.io.load_image(imgname)
    net.blobs['data'].data[...] = transformer.preprocess('data', im)
    out = net.forward()['prob'][0]
    print( '%f %f %d %d' % (out[0],out[1],out.argmax(),label))
